__author__ = 'Yubo Wang'

import torch

def simulate_data(dataset: str = 'kmnist', POS_LABEL: int = 0, ratio: float = 0.3):

    train_dataset = torch.load("processed_data/%s_train_dataset.pth" % dataset)
    test_dataset = torch.load("processed_data/%s_test_dataset.pth" % dataset)

    train_pos_idx = (train_dataset.targets==POS_LABEL)
    test_pos_idx = (test_dataset.targets==POS_LABEL)

    train_data = train_dataset.data[train_pos_idx]
    train_targets = torch.ones(len(train_data))

    test_data = test_dataset.data
    test_data_pos = test_dataset.data[test_pos_idx]
    test_data_neg = test_dataset.data[~test_pos_idx][:int(ratio*len(test_data_pos))]

    test_data = torch.cat((test_data_pos, test_data_neg))
    test_targets = torch.cat((torch.ones(len(test_data_pos)), torch.zeros(len(test_data_neg))))
    
    idx = torch.randperm(len(test_data))
    test_data = test_data[idx]
    test_targets = test_targets[idx]
    
    return train_data, train_targets, test_data, test_targets
